#include<cstdlib>
#include<iostream>
#include<cmath>
#include<algorithm>

#include<formic/genminres.h>
#include<formic/random/random.h>
#include<formic/numeric/numeric.h>
#include<formic/numeric/matrix.h>
#include<formic/lapack/interface.h>
#include<formic/vector/vector.h>
#include<formic/vector/matrix.h>
//#include<formic/test_driver/utilities.h>

namespace formic { namespace genminres_test {

  template<class S, class L> class SimpleMatOp {

    private:

      const formic::Matrix<S,L> & _mat;

    public:

      explicit SimpleMatOp(const formic::Matrix<S,L> & mat) : _mat(mat) {}

      formic::Vector<S,L> operator*(const formic::Vector<S,L> & vec) const {
        const int n = vec.size();
        formic::Vector<S,L> retval(n, 0.0);
        formic::xgemm('N', 'N', n, 1, n, 1.0, &_mat(0,0), n, &vec[0], n, 0.0, &retval[0], n);
        return retval;
      }

  };

  template<class S, class L> class SimpleDiagonalPreconOp {

    private:

      const formic::Matrix<S,L> & _mat;

    public:

      explicit SimpleDiagonalPreconOp(const formic::Matrix<S,L> & mat) : _mat(mat) {}

      formic::Vector<S,L> solve(const formic::Vector<S,L> & vec) const {
        formic::Vector<S,L> retval = vec.clone();
        const double thresh = 1.0e-6;
        for (int i = 0; i < retval.size(); i++)
          if ( _mat(i,i) != 0 )
            retval[i] /= ( std::abs(_mat(i,i)) > thresh ? _mat(i,i) : thresh * _mat(i,i) / std::abs(_mat(i,i)) );
        return retval;
      }

  };

//  template<class S, class L> class VectorNorm {
//
//    public:
//
//      S operator()(const formic::Vector<S,L> & vec) const {
//        S total = 0;
//        for (int i = 0; i < vec.size(); i++) {
//          const S val = std::abs(vec[i]);
//          total += val * val;
//        }
//        return std::sqrt(total);
//      }
//
//  };
//
//  template<class S, class L> class VectorDot {
//
//    public:
//
//      S operator()(const formic::Vector<S,L> & v1, const formic::Vector<S,L> & v2) const {
//        if (v1.size() != v2.size())
//          throw formic::Exception("vectors must be the same length to take their dot product");
//        S total = 0;
//        for (int i = 0; i < v1.size(); i++)
//          total += v1[i] * v2[i];
//        return total;
//      }
//
//  };

} }

int main(int argc, char ** argv) {

  formic::mpi::init(argc, argv);

  // open the output file
  formic::open_output_file("output.txt");

  // set random seed
  formic::set_seed(493754);
  //formic::set_seed(std::time(0));

  // set the matrix dimensions
  const int n = 600;

  // create arrays needed by lapack diagonalization routines and matrix inversion function
  std::vector<double> w(n, 0.0);
  std::vector<double> wr(n, 0.0);
  std::vector<double> wi(n, 0.0);
  int lwork = 10 * n;
  std::vector<double> work(lwork, 0.0);
  std::vector<double> rwork(lwork, 0.0);
  int info = 0;
  std::vector<double> A(n*n, 0.0);
  std::vector<int> iwork(2*n, 0);
  double det = 0.0;

  // build the matrices
  formic::Matrix<double,int> H(n, n, 0.0);
  for (int j = 0; j < n; j++)
  for (int i = 0; i < n; i++) {
    H(i,j) = -10.0 - 2 * formic::random_number<double>();
    if (i != j) {
      H(i,j) = 0.02 * ( formic::random_number<double>() - 0.5 );
    }
  }

  formic::Matrix<double,int> save_H = H.clone();

  // prepare the b vector
  formic::Vector<double,int> b(n, 0.0);
  for (int i = 0; i < n; i++)
    b[i] = 2 * ( formic::random_number<double>() - 0.5 );

  // prepare the solution vector
  formic::Vector<double,int> x(n, 0.0);
  for (int i = 0; i < n; i++)
    x[i] = 2 * ( formic::random_number<double>() - 0.5 );

  // prepare the operators
  formic::genminres_test::SimpleMatOp<double,int> Hop(H);
  formic::genminres_test::SimpleDiagonalPreconOp<double,int> Mop(H);

  // set the solution threshold and maximum iteration number
  const double thresh = 1.0e-8;
  const int max_it = n;

  // set the restart frequency
  const int rf = 3;

  // run the generalized minimal residual solver
  {
    double tol = thresh;
    int m = rf;
    int mi = max_it;
    formic::Matrix<double,int> temp_mat(rf+1, rf+1, 0.0);
    formic::genminres_classes::VectorNorm<double,int> norm;
    formic::genminres_classes::VectorDot<double,int> dot;
    formic::genminres(Hop,
                      x,
                      b,
                      Mop,
                      temp_mat,
                      norm,
                      dot,
                      m,
                      mi,
                      tol);
  }

  // check the solution
  {
    H = save_H.clone();
    formic::Vector<double,int> resid(n, 0.0);
    formic::xgemm('N', 'N', n, 1, n, 1.0, &H(0,0), n, &x[0], n, 0.0, &resid[0], n);
    resid -= b;
    double max_error = 0.0;
    for (int i = 0; i < n; i++)
      max_error = std::max(max_error, std::abs(resid[i]));
    formic::of << boost::format("max error in Ax-b = %.2e") % max_error;
    if (max_error > n * thresh)
      throw formic::Exception("large residual element of %.2e") % max_error;
  }

  formic::close_output_file();

  formic::mpi::finalize();

  return EXIT_SUCCESS;

}
